package com.outertrack.multiproxy;

import org.apache.commons.httpclient.Header;
import org.apache.commons.httpclient.HttpClient;
import org.apache.commons.httpclient.MultiThreadedHttpConnectionManager;
import org.apache.commons.httpclient.methods.GetMethod;
import org.apache.commons.httpclient.methods.HeadMethod;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;

import java.io.*;
import java.util.Arrays;
import java.util.concurrent.*;

/**
 * TODO javadoc
 * <p/>
 * TODOs - BUG when using very small segment sizes
 * <p/>
 * - if executorService.submit() does not block if no threads are available, we
 * need to block a different way to avoid allocating all the buffers upfront
 * <p/>
 * - offer segments to downloadedSegments that are smaller than the
 * availableSegments, so we can start writing the data out sooner
 */
public class MultiThreadDownloader {

  private static final Log log = LogFactory.getLog(MultiThreadDownloader.class);

  // TODO: these should be configurable
  private static final int MAX_CHUNK_SIZE = (1 * 1024);
  private static final int INTERMEDIATE_BUFFER_SIZE = 4096;

  private String url;
  private HttpClient httpClient;
  private static final String CONTENT_LENGTH = "content-length";

  private final BlockingQueue<Segment> availableSegments;
  private final PriorityBlockingQueue<Segment> downloadedSegments;

  private boolean finished;
  private final ExecutorService executorService;
  private static final int MAX_THREADS = 6;

  public static void main(String[] args)
      throws IOException, InterruptedException {

    MultiThreadDownloader downloader = new MultiThreadDownloader(args[0]);

    OutputStream out = new BufferedOutputStream(new FileOutputStream(args[1]));
    downloader.downloadToStream(out);
    System.exit(0);
  }

  public MultiThreadDownloader(String url) {
    this.url = url;

    availableSegments = new ArrayBlockingQueue<Segment>(10);
    downloadedSegments = new PriorityBlockingQueue<Segment>();
    executorService = Executors.newFixedThreadPool(MAX_THREADS);
  }

  public void downloadToStream(OutputStream out)
      throws IOException, InterruptedException {

    httpClient = new HttpClient(new MultiThreadedHttpConnectionManager());

    log.info("Executing HEAD on " + url);
    HeadMethod headMethod = new HeadMethod(url);
    httpClient.executeMethod(headMethod);

    int statusCode = headMethod.getStatusCode();
    if (statusCode >= 400) {
      // TODO handle
      throw new RuntimeException("got status code " + statusCode);
    }

    Header contentLengthHeader = headMethod.getResponseHeader(CONTENT_LENGTH);
    if (contentLengthHeader == null) {
      throw new RuntimeException(
          "missing content-length header not supported yet");
    }

    int contentLength = Integer.parseInt(contentLengthHeader.getValue());
    log.info("Content-length=" + contentLength);

    // TODO: check for Accept-ranges header

    executorService.submit(new SegmentProducer(contentLength));
    executorService.submit(new SegmentConsumer());

    // now wait for segments to come into the downloadedQueue, and write
    // them out to the outputstream, in order since this is a priority queue
    // if the next available segment is not the one we want, then wait
    // until the right one is added
    int offset = 0;
    while (offset < contentLength - 1) {
      Segment segment = downloadedSegments.peek();
      log.debug("Peeked " + segment);
      if (segment == null || segment.getGlobalOffset() != offset) {
        log.debug("Don't want it");
        // TODO: change to use a semaphore or similar thing instead of Thread.wait
        Thread.sleep(1000);
        continue;
      }

      // because this is a priority queue, this segment should be the same
      // one we just polled
      segment = downloadedSegments.take();
      assert offset == segment.getGlobalOffset() :
          "Expected globalOffset " + offset + " but was " +
              segment.getGlobalOffset();

      out.write(segment.getBytes(), 0, segment.getLength());
      offset += segment.getLength();
    }

    out.flush();
    finished = true;
    log.info("All done with download");

  }

  private static class Segment implements Comparable<Segment> {
    private final int globalOffset;
    private final byte[] bytes;

    public Segment(int offset, int length) {
      this.globalOffset = offset;
      bytes = new byte[length];
    }

    public int getGlobalOffset() {
      return globalOffset;
    }

    public int getLength() {
      return bytes.length;
    }

    public byte[] getBytes() {
      //byte[] bytesWritten = new byte[lengthWritten];
      //System.arraycopy(bytes,0,bytesWritten,0,lengthWritten);
      return bytes;
      //return bytesWritten;
    }

    /**
     * Write bytes from the given buffer to this segment. Bytes are copied from
     * the beginning of the given buffer, up to the given length.
     *
     * @param buffer
     * @param internalOffset
     * @param length
     */
    public void write(byte[] buffer, int internalOffset, int length) {
      System.arraycopy(buffer, 0, bytes, internalOffset, length);
    }

    public int compareTo(Segment other) {
      return this.getGlobalOffset() - other.getGlobalOffset();
    }

    public String toString() {
      return "Segment:{globalOffset=" + getGlobalOffset() + ",length=" +
          getLength() + "}";
    }

    public boolean equals(Object o) {
      if (this == o) return true;
      if (o == null || getClass() != o.getClass()) return false;

      Segment that = (Segment) o;

      if (getLength() != that.getLength()) return false;
      if (getGlobalOffset() != that.getGlobalOffset()) return false;
      if (!Arrays.equals(getBytes(), that.getBytes())) return false;

      return true;
    }

    public int hashCode() {
      int result;
      result = getGlobalOffset();
      result = 31 * result + getLength();
      result = 31 * result + Arrays.hashCode(getBytes());
      return result;
    }

  }

  private class SegmentProducer implements Runnable {
    private int contentLength;

    public SegmentProducer(int contentLength) {
      this.contentLength = contentLength;
    }

    public void run() {
      int offset = 0;
      while (offset < contentLength) {
        int length = Math.min(MAX_CHUNK_SIZE, contentLength - offset);
        Segment segment = new Segment(offset, length);
        log.info("Adding " + segment);
        availableSegments.offer(segment);
        offset += length;
      }
      log.info("Done adding segments");
    }
  }

  private class SegmentConsumer implements Runnable {

    public void run() {
      try {
        // get the next available segment to download. if there are no segments
        // in the queue, after waiting up to one second, continue checking.
        // eventually the downloaded will be finished, and we'll stop polling.
        while (!finished) {
          Segment segment = availableSegments.poll(1, TimeUnit.SECONDS);
          if (segment == null) {
            continue;
          }

          log.info("Taking " + segment);
          executorService.submit(new SegmentDownloader(segment));
        }
      } catch (InterruptedException e) {
        Thread.currentThread().interrupt();
        if (!finished) {
          log.error("SegmentConsumer interrupted before finishing!", e);
        }
      }
    }
  }

  private class SegmentDownloader implements Runnable {

    private Segment segment;

    public SegmentDownloader(Segment segment) {
      this.segment = segment;
    }

    public void run() {
      InputStream in = null;
      GetMethod getMethod = null;
      try {
        log.info("Starting to download " + segment);
        getMethod = new GetMethod(url);

        // see http://www.w3.org/Protocols/rfc2616/rfc2616-sec14.html#sec14.35
        // byte offset in Range header are inclusive
        String range =
            "bytes=" + segment.getGlobalOffset() + "-" +
                (segment.getGlobalOffset() + segment.getLength() - 1);
        getMethod.setRequestHeader("Range", range);

        log.info("Executing GET with range " + range);
        httpClient.executeMethod(getMethod);

        in = getMethod.getResponseBodyAsStream();
        byte buf[] = new byte[INTERMEDIATE_BUFFER_SIZE];
        int bytesRead = -1, bufferOffset = 0;
        while ((bytesRead = in.read(buf)) != -1) {
          segment.write(buf, bufferOffset, bytesRead);
          bufferOffset += bytesRead;
        }
        log.info("bufferOffset at end=" + bufferOffset);

        downloadedSegments.offer(segment);

      } catch (Exception e) {
        // for now, log and put back in availableSegments queue to be re-downloaded
        log.warn("Exception downloading segment " + segment, e);
        // TODO have max number of failures, to avoid infinite loop
        availableSegments.offer(segment);
      }
      finally {
        try {
          if (in != null) in.close();
        } catch (IOException e) {
          log.warn("Could not close inputstream", e);
          // not much we can do at this point?
        }
        if (getMethod != null) getMethod.releaseConnection();
      }
    }
  }
}
